import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision.transforms as transforms
import itertools
from PIL import Image
from torchvision import models
from torch.autograd import Variable



class DnSRFD(nn.Module):
    def __init__(self,in_channels=1,mid_channels=256,growth_rate=48):
        super(DnSRFD, self).__init__()

        self.SFE = SFE(in_channels,32,64,mid_channels)

        self.MRDB1 = DenseBlock(6,mid_channels,growth_rate)
        self.Pool1 = Pool(mid_channels,mid_channels*2)
        self.MRDB2 = DenseBlock(6, mid_channels*2, growth_rate)
        self.Pool2 = Pool(mid_channels*2,mid_channels*4)
        self.MRDB3 = DenseBlock(6, mid_channels*4, growth_rate)
        self.Pool3 = Pool(mid_channels*4,mid_channels*8)
        self.MRDB4 = DenseBlock(6, mid_channels*8, growth_rate)
        self.Pool4 = Pool(mid_channels*8,mid_channels*16)
        # Initial convolutional layers



        self.Decovn1 = Decovn(mid_channels * 16, mid_channels * 8)
        self.Down1 = Down(mid_channels * 16)
        self.Decovn2 = Decovn(mid_channels * 8, mid_channels * 4)
        self.Down2 = Down(mid_channels * 8)
        self.Decovn3 = Decovn(mid_channels * 4, mid_channels * 2)
        self.Down3 = Down(mid_channels * 4)
        self.Decovn4 = Decovn(mid_channels * 2, mid_channels)


        self.UPS1 = Decovn(mid_channels, mid_channels)
        self.UPS2 = Decovn(mid_channels, 64)

        self.UPS3 = Outlayer(64,1)


    def forward(self, x):

        out = self.SFE(x)

        out1 = self.Pool1(self.MRDB1(out))
        out2 = self.Pool2(self.MRDB2(out1))
        out3 = self.Pool3(self.MRDB3(out2))
        out4 = self.Pool4(self.MRDB4(out3))

        de_out1 = self.Decovn1(out4)
        de_out1 = self.Down1(torch.cat((out3,de_out1),dim=1))
        de_out2 = self.Decovn2(de_out1)
        de_out2 = self.Down2(torch.cat((out2,de_out2),dim=1))
        de_out3 = self.Decovn3(de_out2)
        de_out3 = self.Down3(torch.cat((out1, de_out3),dim=1))
        de_out3 = self.Decovn4(de_out3)

        out5 = self.UPS1(de_out3)
        out6 = self.UPS2(out5)
        out7 = self.UPS3(out6)

        return out7

class Down(nn.Module):
    def __init__(self, inchannels):
        super(Down, self).__init__()
        self.head = torch.nn.Sequential(
            nn.Conv2d(inchannels, inchannels//2, kernel_size=3, padding=1),
            nn.BatchNorm2d( inchannels//2),
            nn.ELU(inplace=True),
            nn.Conv2d( inchannels//2, inchannels//2, kernel_size=3, padding=1),
            nn.BatchNorm2d( inchannels//2),
            nn.ELU(inplace=True),
        )

    def forward(self, x):
        return self.head(x)

class SFE(nn.Module):
    def __init__(self, channel1, channel2,channel3,channel4):
        super(SFE, self).__init__()
        self.head = torch.nn.Sequential(
            # 浅层提取第一阶段
            torch.nn.Conv2d(in_channels=channel1, out_channels=channel2, kernel_size=3, padding=1),
            torch.nn.BatchNorm2d(channel2),
            torch.nn.ELU(inplace=True),

            torch.nn.MaxPool2d(kernel_size=2, stride=2),
            # 浅层提取第二阶段
            torch.nn.Conv2d(in_channels=channel2, out_channels=channel3, kernel_size=3, padding=1),
            torch.nn.BatchNorm2d(channel3),
            torch.nn.ELU(inplace=True),
            torch.nn.Conv2d(in_channels=channel3, out_channels=channel4, kernel_size=3, padding=1),
            torch.nn.BatchNorm2d(channel4),
            torch.nn.ELU(inplace=True),
            torch.nn.MaxPool2d(kernel_size=2, stride=2)
        )

    def forward(self, x):
        return self.head(x)


# 定义密集块子块-- Bottle-Neck
def conv_bottleneck(in_channels, out_channels):
    p_c1 = int(out_channels // 4)
    p_c2 = int(out_channels // 2)
    p_c3 = int(out_channels // 4)

    block = torch.nn.Sequential(
        Inception(in_c=in_channels,c1=p_c1,c2=[int(p_c2//2),p_c2],c3=[int(p_c3//2),p_c3]),
        torch.nn.BatchNorm2d(out_channels),
        torch.nn.ELU(inplace=True),
    )
    return block

# 定义密集块-- Dense-Block
class DenseBlock(torch.nn.Module):

    def __init__(self, num_layers, in_channels, growth_rate):
        '''
        初始化密集块网络

        :param num_layers: 密集块中bottleneck的层数
        :param in_channels: 当前密集块的输入通道数
        :param growth_rate: 密集块的增长率
        '''
        super(DenseBlock, self).__init__()
        block = []
        channel = in_channels
        # 自定义生成密集块并封装到block中
        for i in range(num_layers):
            block.append(conv_bottleneck(channel, growth_rate))
            channel += growth_rate

        self.covn = torch.nn.Sequential(
            torch.nn.Conv2d(in_channels=channel, out_channels=in_channels, kernel_size=1, stride=1),
            torch.nn.BatchNorm2d(in_channels),
            torch.nn.ELU(inplace=True),
        )

        # 将密集块序列化成为网络的一部分
        self.DenseBlock_Net = torch.nn.ModuleList(block)

    def forward(self, x):
        Y = x
        for layer in self.DenseBlock_Net:
            out = layer(x)
            # 将两个张量拼接在一起（按维数1拼接）
            x = torch.cat((x, out), dim=1)

        return Y + self.covn(x)

class Inception(nn.Module):
    def __init__(self,in_c,c1,c2,c3):
        super(Inception,self).__init__()
        self.p1 = nn.Sequential(
            nn.Conv2d(in_c,c1,kernel_size=1),
            nn.BatchNorm2d(c1),
            nn.ELU(inplace=True)
        )
        self.p2 = nn.Sequential(
            nn.Conv2d(in_c,c2[0],kernel_size=1),
            nn.BatchNorm2d(c2[0]),
            nn.ELU(inplace=True),
            nn.Conv2d(c2[0], c2[1], kernel_size=3,padding=1),
            nn.BatchNorm2d(c2[1]),
            nn.ELU(inplace=True)
        )
        self.p3 = nn.Sequential(
            nn.Conv2d(in_c, c3[0], kernel_size=1),
            nn.BatchNorm2d(c3[0]),
            nn.ELU(inplace=True),
            nn.Conv2d(c3[0], c3[1], kernel_size=3, stride=1,padding=1),
            nn.BatchNorm2d(c3[1]),
            nn.ELU(inplace=True),
            nn.Conv2d(c3[1], c3[1], kernel_size=3, stride=1,padding=1),
            nn.BatchNorm2d(c3[1]),
            nn.ELU(inplace=True)
        )
    def forward(self, x):
        p1 = self.p1(x)
        p2 = self.p2(x)
        p3 = self.p3(x)
        # p4 = self.p4(x)
        return torch.cat((p1,p2,p3),dim=1)


class Pool(nn.Module):
    """Downscaling with maxpool then double conv"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ELU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ELU(inplace=True),
        )

    def forward(self, x):
        return self.maxpool_conv(x)

class Decovn(nn.Module):
    """Downscaling with maxpool then double conv"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.de_covn = nn.Sequential(
            torch.nn.ConvTranspose2d(in_channels=in_channels, out_channels=in_channels, kernel_size=3, stride=2, padding=1,
                                     dilation=1, output_padding=1),
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ELU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ELU(inplace=True),
        )

    def forward(self, x):
        return self.de_covn(x)

class Outlayer(nn.Module):
    """Downscaling with maxpool then double conv"""

    def __init__(self, in_channels,out_channels):
        super().__init__()
        # 设置解释层（输出层）
        self.classifier = torch.nn.Sequential(
            torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1),
            torch.nn.BatchNorm2d(in_channels),
            nn.ELU(inplace=True),
            torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1),
            torch.nn.BatchNorm2d(out_channels),
            nn.ELU(inplace=True),
            torch.nn.Conv2d(out_channels, out_channels, kernel_size=1,stride=1,padding=0),


        )

    def forward(self, x):
        return self.classifier(x)